import numpy as np
import pandas as pd


def z_score(x):
  return (x - x.mean()) / x.std(ddof=0)


def zscore(x, window):
  r = x.rolling(window=window)
  m = r.mean().shift(1)
  s = r.std(ddof=0).shift(1)
  z = (x - m) / s
  return z


def cross_vwap(df, x_itv_s, x_itv_l):
  """cross, X:prod_0_vwap(t-{x_itv_s}m ~ t)-prod_1_vwap(t-{x_itv_l}m ~ t)"""
  x = - (df[f't.vwap.{x_itv_s}m/prod_0'] - df[f't.vwap.{x_itv_l}m/prod_0']) \
      + (df[f't.vwap.{x_itv_s}m/prod_1'] - df[f't.vwap.{x_itv_l}m/prod_1'])
  x = z_score(x)
  # x[(x > 1.5) | (x < -1.5)] = 0
  # x /= 1.5

  x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])

  # x[(x > 15) & (x < 25)] = 0
  # x[(x < -15) & (x > -25)] = 0
  # x[(x > 25) | (x < -25)] = 0
  # x[(x > 35) | (x < -35)] = -x[(x > 35) | (x < -35)]
  # x[(x > 25) | (x < -25)] = -x[(x > 25) | (x < -25)]
  return x


def self_mid_vwap(df, x_itv, roll_itv):
  """X:vwap(t-{x_itv}m ~ t)-mid_p(t)"""
  x = df[f't.vwap.{x_itv}m/prod_0'] - df['b.mid_p/prod_0']
  x = zscore(x, roll_itv)

  x[(x > 2) | (x < -2)] = 0
  x /= 2.0

  # x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def self_volume_sq_diff(df, roll_itv):
  """self, X:ma(volume^2 * (vwap(t-1m ~ t)-vwap(t-2m ~ t-1m)), {roll_itv}min)"""

  base_vwap_sign = df['t.vwap.1m/prod_0'] - df['t.vwap.1m/prod_0'].shift(1)

  base_vwap_sign[base_vwap_sign > 0] = 1
  base_vwap_sign[base_vwap_sign < 0] = -1

  base_panel_vol = df['t.q.sum1m/prod_0'] / df['t.q.sum1m/prod_0'].mean()
  base_volume_sq = base_panel_vol**2

  base_sign_sq = base_vwap_sign * base_volume_sq

  x = -base_sign_sq.rolling(roll_itv).sum()
  # x = x - comp_sign_sq.rolling(interval_mins).sum()
  x = z_score(x)
  x[(x > 1) | (x < -1)] = 0

  return x


def cross_volume_sq_diff(df, roll_itv):
  """cross, X:ma(volume^2 * (vwap(t-1m ~ t)-vwap(t-2m ~ t-1m)), {roll_itv}min)"""

  base_vwap_sign = df['t.vwap.1m/prod_0'] - df['t.vwap.1m/prod_0'].shift(1)
  comp_vwap_sign = df['t.vwap.1m/prod_1'] - df['t.vwap.1m/prod_1'].shift(1)

  base_vwap_sign[base_vwap_sign > 0] = 1
  base_vwap_sign[base_vwap_sign < 0] = -1
  comp_vwap_sign[comp_vwap_sign > 0] = 1
  comp_vwap_sign[comp_vwap_sign < 0] = -1

  base_panel_vol = df['t.q.sum1m/prod_0'] / df['t.q.sum1m/prod_0'].mean()
  comp_panel_vol = df['t.q.sum1m/prod_1'] / df['t.q.sum1m/prod_1'].mean()
  base_volume_sq = base_panel_vol**2
  comp_volume_sq = comp_panel_vol**2

  base_sign_sq = base_vwap_sign * base_volume_sq
  comp_sign_sq = comp_vwap_sign * comp_volume_sq

  x = -base_sign_sq.rolling(roll_itv).sum()
  x = x - comp_sign_sq.rolling(roll_itv).sum()
  x = z_score(x)
  x[(x > 0.5) | (x < -0.5)] = 0
  x *= 2.0
  return x


def vwap_ma_diff(df, x_itv, roll_itv):
  """X:p_x[t] - p_y[t] - ma(p_x - p_y, [t-{x_itv}m, t])"""
  diff = df[f't.vwap.{x_itv}m/prod_0'] - df[f't.vwap.{x_itv}m/prod_1']

  diff_ma = diff.rolling(roll_itv).mean()
  x = diff_ma - diff

  x = z_score(x)
  x[(x > 1.5) | (x < -1.5)] = 0
  x /= 1.5
  return x


def vwap_div_mid_p(df, x_itv):
  """cross, X: prod_0(vwap{x_itv}m / mid_p) - prod_1(vwap{x_itv}m / mid_p)"""
  x = (df[f't.vwap.{x_itv}m/prod_0'] / df['b.mid_p/prod_0'])
  x = x - (df[f't.vwap.{x_itv}m/prod_1'] / df['b.mid_p/prod_1'])
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  # x[(x > 5) | (x < -5)] = 0
  x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def use_bitmex_oi(df):
  """cross, -(prod_0(mid_ret.1m) - prod_1(mid_ret.1m)) / bitmex.oi"""

  # x = -(df['mid_ret.1m/prod_0'] - df['mid_ret.1m/prod_1']) # / df['i.oi']
  # x = x / zscore(df['oi.1m'], 300)
  x = df['oi.10m'] / (df['t.v.sum10m/prod_1'] / df['t.v.sum10m/prod_0'])
  # t = z_score(df['oi.5m'])
  # x = x / t
  # x = df['i.oi']
  x = x.replace([np.inf, -np.inf], np.nan)
  x = (x - x.mean()) / x.std(ddof=0)
  x[(x > 0.5) | (x < -0.5)] = 0
  y = df['Y.mid_ret.1m/prod_0']
  t = df['global_timestamp']
  return x, y, t


def buy_sell_trade_pq_diff(df, is_cross, x_itv):
  """cross:{is_cross}, X: f(prod0) - f(prod1), f=(buy.v.sum{x_itv}m - sell.v.sum{x_itv}m)"""
  x0 = (df[f't.buy.v.msum{x_itv}m/prod_0'] - df[f't.sell.v.msum{x_itv}m/prod_0'])
  x = x0
  if is_cross:
    x1 = (df[f't.buy.v.msum{x_itv}m/prod_1'] - df[f't.sell.v.msum{x_itv}m/prod_1'])
    x = x - x1
  else:
    x = -x

  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  # x[(x > 1.5) | (x < -1.5)] = 0
  x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])

  return x


def buy_sell_trade_q_diff(df, is_cross, x_itv):
  """cross:{is_cross}, X: f(prod0) - f(prod1), f=(buy.q.sum{x_itv}m - sell.q.sum{x_itv}m)"""
  x0 = (df[f't.buy.q.msum{x_itv}m/prod_0'] - df[f't.sell.q.msum{x_itv}m/prod_0'])
  x = x0
  if is_cross:
    x1 = (df[f't.buy.q.msum{x_itv}m/prod_1'] - df[f't.sell.q.msum{x_itv}m/prod_1'])
    x = x - x1
  else:
    x = -x

  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  # x[(x > 2) | (x < -2)] = 0
  x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def mid_p_rev(df, x_itv):
  """self, X: -mid_ret.{x_itv}m"""
  # print(list(df))
  x0 = (df[f'mid_ret.{x_itv}m/prod_0'])
  x = -x0
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  # x[(x > 3) | (x < -3)] = 0
  x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def high_close_low(df, x_itv):
  """self, X: {x_itv}min, (high - close) - (close - low)"""
  x0 = df[f't.p.mmax{x_itv}m/prod_0'] - df['t.p/prod_0']
  x1 = df['t.p/prod_0'] - df[f't.p.mmin{x_itv}m/prod_0']
  x = x0 - x1
  x = z_score(x)
  x[(x > 2) | (x < -2)] = 0
  x /= 2.0
  # x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  # x[(x > 1) | (x < -1)] = 0

  return x


def high_low_vwap(df, x_itv):
  """self, X: {x_itv}min, (((high * low)^0.5) - vwap) """
  x0 = np.sqrt(df[f't.p.mmax{x_itv}m/prod_0'] * df[f't.p.mmin{x_itv}m/prod_0'])
  x1 = df[f't.vwap.{x_itv}m/prod_0']
  x = x0 - x1
  x = z_score(x)
  x[(x > 2) | (x < -2)] = 0
  x /= 2.0
  # x[(x > 1) | (x < -1)] = 0
  # x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def vwap_close(df, x_itv):
  """self, X: {x_itv}min, (vwap-close) / (vwap+close) """
  x0 = df[f't.vwap.{x_itv}m/prod_0'] - df['t.p/prod_0']
  x1 = df[f't.vwap.{x_itv}m/prod_0'] + df['t.p/prod_0']
  x = x0 / x1
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  x[(x > 2) | (x < -2)] = 0
  x /= 2.0
  # x[(x > 1) | (x < -1)] = 0
  # x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def high_avg(df, x_itv):
  """self, X: (((sum(high, {x_itv}min) / {x_itv}min) < high) ? (-1 * delta(high, 2)) : 0)"""
  m = df['t.p.mmax1m/prod_0'].rolling(window=x_itv, min_periods=1).mean()
  bef_high = df['t.p.mmax1m/prod_0'].shift(2)
  adf = pd.DataFrame({'ma': m, 'bef_high': bef_high, 'high': df['t.p.mmax1m/prod_0']})

  # print(adf)

  def comp(ele):
    if ele.ma < ele.high:
      return -(ele.high - ele.bef_high)
    return np.nan

  x = adf.apply(comp, axis=1)
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x).fillna(0)
  x[(x > 2) | (x < -2)] = 0
  x /= 2.0
  # x[(x > 1) | (x < -1)] = np.sign(x[(x > 1) | (x < -1)])
  return x


def low_avg(df, x_itv):
  """self, X: (((sum(min, {x_itv}min) / {x_itv}min) > min) ? (- delta(min), 2)) : 0)"""
  m = df['t.p.mmin1m/prod_0'].rolling(window=x_itv, min_periods=1).mean()
  bef_low = df['t.p.mmin1m/prod_0'].shift(2)
  adf = pd.DataFrame({'ma': m, 'bef_low': bef_low, 'low': df['t.p.mmin1m/prod_0']})

  # print(adf)

  def comp(ele):
    if ele.ma > ele.low:
      return -(ele.low - ele.bef_low)
    return np.nan

  x = adf.apply(comp, axis=1)
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x).fillna(0)
  x[(x > 2) | (x < -2)] = 0
  x /= 2.0
  return x


def mix_high_low(df, x_itv):
  """self, X: high_avg + low_avg, {x_itv}min"""
  return high_avg(df, x_itv) + low_avg(df, x_itv)


############## For test ###############


def book_pressure_trade_q_diff(df, x_itv, y_itv):
  """self, X: f=(t.q.sum{x_itv}m / (b.bid0_q + b.ask0_q))
     Y:vwap_ret.{y_itv}m"""
  x0 = (df[f't.q.sum{x_itv}m/prod_0']) / \
    (df['b.bid0_q/prod_0'] + df['b.ask0_q/prod_0'])

  x = x0
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  x[(x > 0.5) | (x < -0.5)] = 0
  y = df[f'Y.vwap_ret.{y_itv}m/prod_0']
  t = df['global_timestamp']
  return x, y, t


def bp_diff(df, x_itv, y_itv):
  """self, X: f=(b.bid0_q - b.ask0_q))
     Y:vwap_ret.{y_itv}m"""
  # print(list(df))
  x0 = (df['b.bid0_q.msum5m/prod_0'] - df['b.ask0_q.msum5m/prod_0'])
  x = x0
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  x[(x > 0.5) | (x < -0.5)] = 0
  y = df[f'Y.vwap_ret.{y_itv}m/prod_0']
  t = df['global_timestamp']
  return x, y, t


def mid_p_spread(df, x_itv, y_itv):
  """self, X: mid_ret.{x_itv}m
     Y:vwap_ret.{y_itv}m"""
  # print(list(df))
  x0 = df['t.v.']
  x = x0
  x = x.replace([np.inf, -np.inf], np.nan)
  x = z_score(x)
  x[(x > 3) | (x < -3)] = 0
  y = df[f'Y.vwap_ret.{y_itv}m/prod_0']
  t = df['global_timestamp']
  return x, y, t


"""
['global_timestamp', 'b.ask0_p/prod_0', 'b.ask0_q/prod_0', 
'b.bid0_p/prod_0', 'b.bid0_q/prod_0', 'b.spread/prod_0', 
'b.mid_p/prod_0', 'b.mid_p.mstd1m/prod_0', 'b.mid_p.mstd5m/prod_0',
'b.mid_p.mstd10m/prod_0', 'b.mid_p.mstd20m/prod_0',
'b.mid_p.mstd30m/prod_0', 'b.mid_p.mstd1h/prod_0',
'mid_ret.1m/prod_0', 'mid_ret.5m/prod_0', 'mid_ret.10m/prod_0',
'mid_ret.20m/prod_0', 'mid_ret.30m/prod_0', 'mid_ret.60m/prod_0',
't.vwap.1m/prod_0', 't.vwap.5m/prod_0', 't.vwap.10m/prod_0',
't.vwap.20m/prod_0', 't.vwap.30m/prod_0', 't.vwap.60m/prod_0',
't.q.sum1m/prod_0', 't.q.sum5m/prod_0', 't.q.sum10m/prod_0', 
't.q.sum20m/prod_0', 't.q.sum30m/prod_0', 't.q.sum60m/prod_0', 
't.v.sum1m/prod_0', 't.v.sum5m/prod_0', 't.v.sum10m/prod_0', 
't.v.sum20m/prod_0', 't.v.sum30m/prod_0', 't.v.sum60m/prod_0',
't.p.mstd1m/prod_0', 't.p.mstd5m/prod_0', 't.p.mstd10m/prod_0', 
't.p.mstd20m/prod_0', 't.p.mstd30m/prod_0', 't.p.mstd1h/prod_0',
't.buy.q.msum1m/prod_0', 't.buy.q.msum5m/prod_0', 't.buy.q.msum10m/prod_0', 
't.buy.q.msum20m/prod_0', 't.buy.q.msum30m/prod_0', 't.buy.q.msum1h/prod_0', 
't.buy.v.msum1m/prod_0', 't.buy.v.msum5m/prod_0', 't.buy.v.msum10m/prod_0',
't.buy.v.msum20m/prod_0', 't.buy.v.msum30m/prod_0', 't.buy.v.msum1h/prod_0',
't.sell.q.msum1m/prod_0', 't.sell.q.msum5m/prod_0', 't.sell.q.msum10m/prod_0',
't.sell.q.msum20m/prod_0', 't.sell.q.msum30m/prod_0', 't.sell.q.msum1h/prod_0',
't.sell.v.msum1m/prod_0', 't.sell.v.msum5m/prod_0', 't.sell.v.msum10m/prod_0',
't.sell.v.msum20m/prod_0', 't.sell.v.msum30m/prod_0', 't.sell.v.msum1h/prod_0',
'Y.mid_ret.1m/prod_0', 'Y.mid_ret.5m/prod_0', 'Y.mid_ret.10m/prod_0',
'Y.mid_ret.20m/prod_0', 'Y.mid_ret.30m/prod_0', 'Y.mid_ret.60m/prod_0',
]

"""


def new_self_mid_p_rev(row, products, x_itv):
  """self, X: -mid_ret.{x_itv}m"""
  self_prod = products[0]
  x0 = (row[f'{self_prod}/mid_ret.{x_itv}m'])
  x = -x0
  return x


def new_self_high_close_low(row, products, x_itv):
  """self, X: {x_itv}min, (high - close) - (close - low)"""
  self_prod = products[0]
  x0 = row[f'{self_prod}/t.p.mmax{x_itv}m'] - row[f'{self_prod}/t.p']
  x1 = row[f'{self_prod}/t.p'] - row[f'{self_prod}/t.p.mmin{x_itv}m']
  return x0 - x1


def new_self_vwap_close(row, products, x_itv):
  """self, X: {x_itv}min, (vwap-close) / (vwap+close) """
  self_prod = products[0]
  x0 = row[f'{self_prod}/t.vwap.{x_itv}m'] - row[f'{self_prod}/t.p']
  x1 = row[f'{self_prod}/t.vwap.{x_itv}m'] + row[f'{self_prod}/t.p']
  return x0 / x1


def new_cross_vwap_div_mid_p(row, products, x_itv):
  """cross, X: prod_0(vwap{x_itv}m / mid_p) - prod_1(vwap{x_itv}m / mid_p)"""
  x = (row[f'{products[0]}/t.vwap.{x_itv}m'] / row[f'{products[0]}/b.mid_p'])
  x = x - (row[f'{products[1]}/t.vwap.{x_itv}m'] / row[f'{products[1]}/b.mid_p'])
  return x


def new_self_mid_vwap(row, products, x_itv):
  """X:vwap(t-{x_itv}m ~ t)-mid_p(t)"""
  self_prod = products[0]
  x = row[f'{self_prod}/t.vwap.{x_itv}m'] - row[f'{self_prod}/b.mid_p']
  return x


def new_cross_vwap(row, products, x_itv_s, x_itv_l):
  """cross, X:prod_0_vwap(t-{x_itv_s}m ~ t)-prod_1_vwap(t-{x_itv_l}m ~ t)"""
  x = - (row[f'{products[0]}/t.vwap.{x_itv_s}m'] - row[f'{products[0]}/t.vwap.{x_itv_l}m']) \
      + (row[f'{products[1]}/t.vwap.{x_itv_s}m'] - row[f'{products[1]}/t.vwap.{x_itv_l}m'])
  return x


def new_cross_buy_sell_trade_pq_diff(row, products, x_itv):
  """cross, X: f(prod0) - f(prod1), f=(buy.v.sum{x_itv}m - sell.v.sum{x_itv}m)"""
  x0 = (row[f'{products[0]}/t.buy.v.msum{x_itv}m'] - row[f'{products[0]}/t.sell.v.msum{x_itv}m'])
  x1 = (row[f'{products[1]}/t.buy.v.msum{x_itv}m'] - row[f'{products[1]}/t.sell.v.msum{x_itv}m'])
  return x0 - x1


def new_self_buy_sell_trade_pq_diff(row, products, x_itv):
  """self, X: -f(prod0), f=(buy.v.sum{x_itv}m - sell.v.sum{x_itv}m)"""
  x0 = -(row[f'{products[0]}/t.buy.v.msum{x_itv}m'] - row[f'{products[0]}/t.sell.v.msum{x_itv}m'])
  return x0


def new_cross_buy_sell_trade_q_diff(row, products, x_itv):
  """cross, X: f(prod0) - f(prod1), f=(buy.q.sum{x_itv}m - sell.q.sum{x_itv}m)"""
  x0 = (row[f'{products[0]}/t.buy.q.msum{x_itv}m'] - row[f'{products[0]}/t.sell.q.msum{x_itv}m'])
  x1 = (row[f'{products[1]}/t.buy.q.msum{x_itv}m'] - row[f'{products[1]}/t.sell.q.msum{x_itv}m'])
  return x0 - x1


def new_self_buy_sell_trade_q_diff(row, products, x_itv):
  """self, X: -f(prod0), f=(buy.q.sum{x_itv}m - sell.q.sum{x_itv}m)"""
  x0 = -(row[f'{products[0]}/t.buy.q.msum{x_itv}m'] - row[f'{products[0]}/t.sell.q.msum{x_itv}m'])
  return x0


def new_self_vwap_ma_diff(row, products, x_itv):
  """self, X:p_x[t] - p_y[t] - ma(p_x - p_y, [t-{x_itv}m, t])"""
  diff = (row[f'{products[0]}/t.p'] - row[f'{products[0]}/t.vwap.{x_itv}m'])

  if x_itv == 60:
    div = row[f'{products[0]}/t.p.mstd1h']
  else:
    div = row[f'{products[0]}/t.p.mstd{x_itv}m']
  x = -diff / div
  return x


########################
### this is for test.###
########################
def alpha_1(row, mul):
  from scipy.stats import logistic
  return sum(logistic.cdf(row * mul))
  return x


def alpha_2(row, pow):
  return sum(row**pow)


def temp_trans(row, offset):
  return row + offset
